import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt


class StudentAgent(object):
    def __init__(self, type):
        self.num_neighbors = 12
        self.require_same_type_num = 8
        self.type = type
        self.show_position()

    def show_position(self):
        self.position = np.random.uniform(0, 1), np.random.uniform(0, 1)

    def get_distance(self, other):
        """
        欧式距离
        :param other:
        :return:
        """
        a = self.position[0] - other.position[0]
        b = self.position[1] - other.position[1]
        return np.sqrt(np.power(a, 2) + np.power(b, 2))

    def happy(self, agents):
        distances = []
        for agent in agents:
            if self != agent:
                distance = self.get_distance(agent)
                distances.append((distance, agent))
        # 获取最近的邻居
        distances.sort()
        neighbors = [agent for d, agent in distances[: self.num_neighbors]]
        # 获取与自己相同类型的邻居
        same_type_arr = [self.type == agent.type for agent in neighbors]
        return sum(same_type_arr) >= self.require_same_type_num

    def update(self, agents):
        """
        不幸福时换位置，知道满意为止
        :param agent:
        :return:
        """
        while self.happy(agents):
            # 获取随机位置
            self.show_position()


def plot_distribution(agents, cycle_num):
    """
    画图
    :param agents:
    :param cycle_num:
    :return:
    """
    # matplotlib中文显示方块
    mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
    mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    x1, y1 = [], []
    x2, y2 = [], []
    x3, y3 = [], []

    for agent in agents:
        x, y = agent.position
        if agent.type == 0:
            x1.append(x)
            y1.append(y)
        elif agent.type == 1:
            x2.append(x)
            y2.append(y)
        else:
            x3.append(x)
            y3.append(y)

    plot_args = dict(markersize=8, alpha=0.65)
    plt.plot(x1, y1, 'o', color='r', **plot_args)
    plt.plot(x2, y2, 'o', color='b', **plot_args)
    plt.plot(x3, y3, 'D', color='g', **plot_args)
    plt.title('迭代次数：{}'.format(cycle_num))
    plt.show()


def schelling():
    """
    schelling分离模型（阈值模型）
    场景：
    学生分类：优秀，良好，合格-->0, 1, 2
    每类250名学生
    学生位置(x, y)  0 < x, y < 1
    目标：某学生周围（欧氏距离最近）的12个同学中有超过一半跟他/她一类型，则他/她幸福
    条件：位置相互独立，服从二元均匀分布
    """
    num = 250
    np.random.seed(10)

    agents = [StudentAgent(0) for _ in range(num)]
    agents += [StudentAgent(1) for _ in range(num)]
    agents += [StudentAgent(2) for _ in range(num)]

    count = 0
    # turn = True
    # while turn:
    #     count += 1
    #     no_more_remove = True
    #     for agent in agents:
    #         if not agent.happy(agents):
    #             old_position = agent.position
    #             agent.update(agents)
    #             if old_position != agent.position:
    #                 no_more_remove = False
    #     if no_more_remove:
    #         turn = False
    #     print('count: {}'.format(count))
    #
    # plot_distribution(agents, cycle_num=count)

    terminate = False
    while terminate == False:
        print('count: {}'.format(count))
        plot_distribution(agents, cycle_num=count)
        count += 1
        no_more_remove = True
        for agent in agents:
            old_postion = agent.position
            agent.update(agents)
            if old_postion != agent.position:
                no_more_remove = False
        if no_more_remove:
            terminate = True


schelling()
